#pragma once

#include <iostream>
#include <random>
#include <atomic>

#ifdef USE_MPI
#include "mpi.h"
#endif

#include "timer.hpp"

class BaseTester {
 public:
  BaseTester() {};
  virtual ~BaseTester() {};
  // data generator
  virtual void DataGenerate() = 0;
  // the base sequential implementation
  virtual void Baseline() = 0;
  // the parallel implementation
  virtual void Implementation() = 0;
  // used to compare the result if right
  virtual bool CheckResult() = 0;
  // run for test
  virtual void Run() {
    int rank = 0;
#ifdef USE_MPI
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
#endif
    // 1. gererate the input data
    if (rank == 0) {
      DataGenerate();
    }
    Timer t1;
    // 2. run the base line
    if (rank == 0) {
      Baseline();
    }
    double baseline_time = t1.Passed();
    Timer t2;
    // 3. run the user's implementation
    Implementation();
    double impl_time = t2.Passed();
    // 4. check the result
    bool compare = false;
    if (rank == 0) {
      compare = CheckResult();
      if (compare == false) {
        std::cerr << "The result is \033[31merror\033[0m,"
          << " please make sure the implementation is right" << std::endl;
        exit(1);
      }
      // 5. give the speedup ratio
      std::cout << "Congratulations, test \033[36mPASSED\033[0m. " << std::endl
          << "Your program is \033[33m" << baseline_time / impl_time
          << "X\033[0m faster than baseline." << std::endl;
    }
  }

  virtual void CheckOpenMP() {
    std::atomic<int> count(0);
#pragma omp parallel
    {
      count++;
    }
    if (count == 1) {
      std::cerr << "OpenMP is not working, "
        << "please check your OpenMP setting" << std::endl;
      exit(1);
    }
    else {
      std::cout << "OpenMP is open with \033[33m" << count << "\033[0m threads" << std::endl;
    }
  }

  virtual void CheckMPI() {
#ifdef USE_MPI
    MPI_Comm comm = MPI_COMM_WORLD;
    int size = 0;
    MPI_Comm_size(comm, &size);
    std::cout << "MPI is open with \033[33m" << size << "\033[0m processes" << std::endl;
#endif
  }
};
